
from manimlib import*
import numpy as np

default_color_list=[RED_A,RED_B,RED_C,RED_D,RED_E]

square=Polygon(UR,UL,DL,DR,fill_opacity=1,stroke_width=0,color=WHITE)

class mandobrot_set(VGroup):
    CONFIG={
        "gradient":default_color_list
    }

    def __init__(self, resolution, iteration, **kwargs):
        self.resolution = resolution
        self.iteration = iteration
        super().__init__(*self.get_every_pixel(),**kwargs)
        
    def get_every_pixel(self):
        list=[]
        for i in np.arange(-2,2,4/self.resolution):
            for j in np.arange(-2,2,4/self.resolution):
                list.append(Polygon(np.array([i,j,0]),
                np.array([i,j+4/self.resolution,0]),
                np.array([i+4/self.resolution,j+4/self.resolution,0]),
                np.array([i+4/self.resolution,j,0]),fill_opacity=1,stroke_width=0,color=WHITE))

                print(len(list))
        return list
    
    def set_gradient_color(self):
        for i in range(self.resolution**2):
            coor=self[i].get_center()
            t=0
            start=complex(coor[0],coor[1])
            ori=complex(0,0)
            while t<=self.iteration and abs(ori)<=2:
                ori=ori*ori+start
                t+=1
            if t>=self.iteration:
                color=WHITE
            else:
                color=self.gradient[t % len(self.gradient)]
            self[i].set_color(color)
            print(i)
        return self

class gradient(Scene):
    def construct(self):
        mando=mandobrot_set(200,100).set_gradient_color()
        self.add(mando)
